
import indexing.featureindex
import coverage.basecountreader
import sys

def getPhastCons(phastConsFormat, snpFile):
    snps = {}
    for line in open(snpFile):
        cols = line.strip().split("\t")
        (chr, pos) = cols[0].split(":")
        if not snps.has_key(chr):
            snps[chr] = []
        snps[chr].append(long(pos))
    chrs = snps.keys()[:]
    chrs.sort()
    for chr in chrs:
        print >> sys.stderr, "# Read %d SNPs in %s from file." % (len(snps[chr]), chr)
        snps[chr].sort()
    for chr in chrs:
        print >> sys.stderr, "# Reading %d features from %s" % (len(snps[chr]), phastConsFormat % chr)
        index = indexing.featureindex.FeatureIndex(phastConsFormat % chr)
        reader = coverage.basecountreader.BaseCountReader(index)
        for pos in snps[chr]:
            feature = reader[(chr, pos)]
            if not feature:
                print "%s:%ld\tNA" % (chr, pos)
            else:
                print "%s:%ld\t%s" % (chr, pos, feature.content.strip())

def getPhastConsStreaming(phastConsFormat, snpFile):
    currChr = None
    for line in open(snpFile):
        cols = line.strip().split("\t")
        (chr, pos) = cols[0].split(":")
        pos = long(pos)
        if not currChr or currChr != chr:
            index = indexing.featureindex.FeatureIndex(phastConsFormat % chr)
            reader = coverage.basecountreader.BaseCountReader(index)
            currChr = chr
        feature = reader[(chr, pos)]
        if not feature:
            print "%s:%ld\tNA" % (chr, pos)
        else:
            print "%s:%ld\t%s" % (chr, pos, feature.content.strip())

if __name__=="__main__":
    getPhastConsStreaming("/archive10/bowtie/references/phastcons/%s.phastcons.txt.INDEX", sys.argv[1])

